package com.frp.client;

import com.frp.config.FrpcConfig;
import com.frp.protocol.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.Set;
import java.util.Iterator;

/**
 * 代理处理器
 * 处理单个代理的数据转发
 */
public class ProxyHandler {
    private static final Logger logger = LoggerFactory.getLogger(ProxyHandler.class);
    private static final int BUFFER_SIZE = 8192;
    
    private final FrpcConfig.ProxyConfig proxyConfig;
    private final ConnectionManager connectionManager;
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final ConcurrentHashMap<String, WorkConnection> workConnections = new ConcurrentHashMap<>();
    
    public ProxyHandler(FrpcConfig.ProxyConfig proxyConfig, ConnectionManager connectionManager) {
        this.proxyConfig = proxyConfig;
        this.connectionManager = connectionManager;
    }
    
    /**
     * 启动代理处理器
     */
    public void start() {
        if (running.compareAndSet(false, true)) {
            logger.info("启动代理处理器: {}", proxyConfig.getName());
        }
    }
    
    /**
     * 停止代理处理器
     */
    public void stop() {
        if (running.compareAndSet(true, false)) {
            logger.info("停止代理处理器: {}", proxyConfig.getName());
            
            // 关闭所有工作连接
            for (WorkConnection conn : workConnections.values()) {
                try {
                    conn.close();
                } catch (IOException e) {
                    logger.error("关闭工作连接时出错", e);
                }
            }
            workConnections.clear();
        }
    }
    
    /**
     * 创建工作连接
     */
    public void createWorkConnection() throws IOException {
        SocketChannel workChannel = SocketChannel.open();
        workChannel.configureBlocking(false);
        workChannel.connect(new InetSocketAddress(
            connectionManager.getConfig().getCommon().getServerAddr(), 
            connectionManager.getConfig().getCommon().getServerPort()
        ));
        
        while (!workChannel.finishConnect()) {
            Thread.yield();
        }
        
        logger.info("工作连接已建立");
    }
    public void createWorkConnection(NewWorkConnMessage message) {
        Thread workThread = new Thread(() -> {
            try {
                WorkConnection workConn = new WorkConnection(message);
                workConn.start();
                
                String connectionId = generateConnectionId();
                workConnections.put(connectionId, workConn);
                
                logger.debug("工作连接已创建: {} - {}", proxyConfig.getName(), connectionId);
                
            } catch (IOException e) {
                logger.error("创建工作连接失败", e);
            }
        }, "Work-Connection-" + proxyConfig.getName());
        
        workThread.setDaemon(true);
        workThread.start();
    }
    
    /**
     * 生成连接ID
     */
    private String generateConnectionId() {
        return proxyConfig.getName() + "-" + System.currentTimeMillis();
    }
    
    /**
     * 工作连接内部类
     */
    private class WorkConnection {
        private final NewWorkConnMessage message;
        private SocketChannel serverChannel;
        private SocketChannel localChannel;
        private Selector selector;
        
        public WorkConnection(NewWorkConnMessage message) {
            this.message = message;
        }
        
        /**
         * 启动工作连接
         */
        public void start() throws IOException {
            try {
                // 连接到FRP服务器的工作端口
                serverChannel = SocketChannel.open();
                serverChannel.configureBlocking(false);
                
                InetSocketAddress serverAddress = new InetSocketAddress(
                    connectionManager.getConfig().getCommon().getServerAddr(),
                    connectionManager.getConfig().getCommon().getServerPort()
                );
                
                boolean connected = serverChannel.connect(serverAddress);
                if (!connected) {
                    while (!serverChannel.finishConnect()) {
                        Thread.sleep(10);
                    }
                }
                
                // 发送工作连接响应
                NewWorkConnRespMessage resp = new NewWorkConnRespMessage(
                    message.getProxyName(),
                    null  // 成功
                );
                
                sendMessage(resp);
                
                // 连接到本地服务
                localChannel = SocketChannel.open();
                localChannel.configureBlocking(false);
                
                InetSocketAddress localAddress = new InetSocketAddress(
                    proxyConfig.getLocalIp(),
                    proxyConfig.getLocalPort()
                );
                
                connected = localChannel.connect(localAddress);
                if (!connected) {
                    while (!localChannel.finishConnect()) {
                        Thread.sleep(10);
                    }
                }
                
                // 开始数据转发
                startDataForwarding();
                
            } catch (Exception e) {
                logger.error("工作连接启动失败", e);
                close();
                throw new IOException("工作连接启动失败", e);
            }
        }
        
        /**
         * 发送消息
         */
        private void sendMessage(Message message) throws IOException {
            byte[] data = MessageCodec.encode(message);
            ByteBuffer buffer = ByteBuffer.wrap(data);
            
            while (buffer.hasRemaining()) {
                serverChannel.write(buffer);
            }
        }
        
        /**
         * 开始数据转发
         */
        private void startDataForwarding() throws IOException {
            selector = Selector.open();
            serverChannel.register(selector, SelectionKey.OP_READ);
            localChannel.register(selector, SelectionKey.OP_READ);
            
            logger.debug("开始数据转发: {} -> {}:{}", 
                proxyConfig.getName(),
                proxyConfig.getLocalIp(),
                proxyConfig.getLocalPort()
            );
            
            ByteBuffer serverBuffer = ByteBuffer.allocate(BUFFER_SIZE);
            ByteBuffer localBuffer = ByteBuffer.allocate(BUFFER_SIZE);
            
            while (running.get() && 
                   serverChannel.isOpen() && 
                   localChannel.isOpen()) {
                
                if (selector.select(1000) > 0) {
                    Set<SelectionKey> selectedKeys = selector.selectedKeys();
                    
                    for (SelectionKey key : selectedKeys) {
                        try {
                            if (key.isReadable()) {
                                SocketChannel channel = (SocketChannel) key.channel();
                                
                                if (channel == serverChannel) {
                                    // 从服务器读取数据，转发到本地
                                    int bytesRead = serverChannel.read(serverBuffer);
                                    if (bytesRead == -1) {
                                        close();
                                        return;
                                    }
                                    
                                    if (bytesRead > 0) {
                                        serverBuffer.flip();
                                        while (serverBuffer.hasRemaining()) {
                                            localChannel.write(serverBuffer);
                                        }
                                        serverBuffer.clear();
                                    }
                                    
                                } else if (channel == localChannel) {
                                    // 从本地读取数据，转发到服务器
                                    int bytesRead = localChannel.read(localBuffer);
                                    if (bytesRead == -1) {
                                        close();
                                        return;
                                    }
                                    
                                    if (bytesRead > 0) {
                                        localBuffer.flip();
                                        
                                        // 创建数据消息
                                        byte[] data = new byte[localBuffer.remaining()];
                                        localBuffer.get(data);
                                        
                                        DataMessage dataMessage = new DataMessage(
                                            proxyConfig.getName(),
                                            data
                                        );
                                        
                                        sendMessage(dataMessage);
                                        localBuffer.clear();
                                    }
                                }
                            }
                        } catch (IOException e) {
                            logger.error("数据转发时出错", e);
                            close();
                            return;
                        }
                    }
                    
                    selectedKeys.clear();
                }
            }
        }
        
        /**
         * 关闭连接
         */
        public void close() throws IOException {
            if (selector != null) {
                selector.close();
            }
            if (serverChannel != null) {
                serverChannel.close();
            }
            if (localChannel != null) {
                localChannel.close();
            }
            
            logger.debug("工作连接已关闭: {}", proxyConfig.getName());
        }
    }
}